#include "armgen.h"

#include <cmath>
#include <functional>
#include <iostream>
#include <list>
#include <ostream>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>

#include "arithmetic.h"
#include "armblock.h"
#include "armfunction.h"
#include "arminst.h"
#include "ir.h"
#include "irblock.h"
#include "irfunction.h"
#include "iroperand.h"
#include "memobj.h"
#include "semantics.h"
#include "symbol.h"
#include "types.h"
#include "utils.h"

using Register = Arm::Register;

using namespace Utils;

ArmFunction* ArmGenerator::Build(IRFunction* func) {
    auto armFunc = new ArmFunction();
    armFunc->name = func->name;
    armFunc->entry = new ArmBlock();
    armFunc->exit = new ArmBlock();

    std::unordered_map<IRBlock*, ArmBlock*> blocks;
    std::function<void(IRBlock*)> dfs = [&](IRBlock* cur) {
        if (blocks.count(cur)) return;
        blocks[cur] = new ArmBlock();
        for (auto v : cur->GetNext()) dfs(v);
    };
    dfs(func->entry);

    // Alloc stack memory
    int stackSize = 0;
    std::unordered_map<std::string, int> varOffset;
    for (int i = 0; i < func->args.size(); ++i) {
        varOffset[func->args[i]->name] = -(i * 4 + (i < 4 ? 0 : 8));
    }
    auto allocLocalVar = [&](IR::Variable* var, int size) {
        if (!varOffset.count(var->name)) {
            stackSize += size * 4;
            varOffset[var->name] = stackSize;
        }
        return varOffset[var->name];
    };
    for (auto [blk, _] : blocks) {
        for (auto irBase : blk->irList) {
            if (auto ir = irBase->As<IR::Alloc>()) {
                allocLocalVar(ir->dest, ir->size);
            }
        }
    }

    const int directAccessBound = 256;
    auto loadVar = [&](std::list<Arm::Inst*>& ls, Register reg, IR::Operand* op,
                       Register tmp) {
        if (auto var = op->As<IR::Variable>()) {
            if (var->isGlobal) {
                ls.push_back(new Arm::LoadAddr(reg, var->name));
                if (!var->type.isPtr)
                    ls.push_back(
                        new Arm::Load(reg, new Arm::ImmOffset(reg, 0)));
            } else {
                if (!varOffset.count(var->name)) {
                    allocLocalVar(var, 1);
                }
                int off = -varOffset[var->name];
                Arm::MemObj* mem;
                if (std::abs(off) < directAccessBound) {
                    mem = new Arm::ImmOffset(Register::fp, off);
                } else {
                    ls.push_back(new Arm::LoadImm(reg, off));
                    mem = new Arm::RegOffset(Register::fp, reg, 0);
                }
                ls.push_back(new Arm::Load(reg, mem));
            }
        } else if (auto imm = op->As<IR::Immediate>()) {
            if (reg.type == Type::Basic::Int) {
                ls.push_back(new Arm::LoadImm(reg, imm->value.v.i));
            } else {
                ls.push_back(new Arm::LoadImm(tmp, imm->value.v.i));
                ls.push_back(new Arm::Mov(reg, tmp));
            }
        }
    };
    auto storeVar = [&](std::list<Arm::Inst*>& ls, Register reg,
                        IR::Variable* var, Register tmp) {
        if (var->isGlobal) {
            Assert(!var->type.isPtr, "cannot assign global ptr");
            ls.push_back(new Arm::LoadAddr(tmp, var->name));
            ls.push_back(new Arm::Store(reg, new Arm::ImmOffset(tmp, 0)));
        } else {
            if (!varOffset.count(var->name)) {
                allocLocalVar(var, 1);
            }
            int off = -varOffset[var->name];
            Arm::MemObj* mem;
            if (std::abs(off) < directAccessBound) {
                mem = new Arm::ImmOffset(Register::fp, off);
            } else {
                ls.push_back(new Arm::LoadImm(tmp, off));
                mem = new Arm::RegOffset(Register::fp, tmp, 0);
            }
            ls.push_back(new Arm::Store(reg, mem));
        }
    };
    auto addVar = [&](std::list<Arm::Inst*>& ls, Register rd, Register rs,
                      int imm, Register temp) {
        if (imm == 0) return;
        if (0 < imm && imm < 4096) {
            ls.push_back(new Arm::BinaryOpImm(rd, rs, imm, ArithOp::Add));
        } else if (-4096 < imm && imm < 0) {
            ls.push_back(new Arm::BinaryOpImm(rd, rs, -imm, ArithOp::Sub));
        } else if (imm > 0) {
            ls.push_back(new Arm::LoadImm(temp, imm));
            ls.push_back(new Arm::BinaryOp(rd, rs, temp, ArithOp::Add));
        } else if (imm < 0) {
            ls.push_back(new Arm::LoadImm(temp, -imm));
            ls.push_back(new Arm::BinaryOp(rd, rs, temp, ArithOp::Sub));
        }
    };
    auto genBlock = [&](ArmBlock* blk, std::list<Arm::Inst*>& ls,
                        IRBlock* irs) {
        for (auto irBase : irs->irList) {
            if (auto ir = irBase->As<IR::Operation>()) {
                loadVar(ls, 0, ir->rhs, 5);
                if (ir->op == ArithOp::IntCast ||
                    ir->op == ArithOp::FloatCast) {
                    ls.push_back(new Arm::Mov(128, 0));
                    ls.push_back(new Arm::UnaryOp(128, 128, ir->op));
                    storeVar(ls, 128, ir->dest, 5);
                } else if (ir->dest->type.isPtr ||
                           ir->dest->type.base == Type::Basic::Int) {
                    if (ir->lhs) {
                        loadVar(ls, 1, ir->lhs, 5);
                        ls.push_back(new Arm::BinaryOp(2, 1, 0, ir->op));
                    } else {
                        ls.push_back(new Arm::UnaryOp(2, 0, ir->op));
                    }
                    storeVar(ls, 2, ir->dest, 5);
                } else {
                    ls.push_back(new Arm::Mov(128, 0));
                    if (ir->lhs) {
                        loadVar(ls, 1, ir->lhs, 5);
                        ls.push_back(new Arm::Mov(129, 1));
                        ls.push_back(new Arm::BinaryOp(130, 129, 128, ir->op));
                    } else {
                        ls.push_back(new Arm::UnaryOp(130, 128, ir->op));
                    }
                    storeVar(ls, 130, ir->dest, 5);
                }
            } else if (auto ir = irBase->As<IR::Load>()) {
                loadVar(ls, 0, ir->base, 5);
                loadVar(ls, 1, ir->offset, 5);
                auto mem = new Arm::RegOffset(0, 1, 2);
                int dest = ir->dest->type.base == Type::Basic::Float ? 128 : 2;
                ls.push_back(new Arm::Load(dest, mem));
                storeVar(ls, dest, ir->dest, 5);
            } else if (auto ir = irBase->As<IR::Store>()) {
                loadVar(ls, 0, ir->base, 5);
                loadVar(ls, 1, ir->offset, 5);
                auto mem = new Arm::RegOffset(0, 1, 2);
                int src = ir->src->type.base == Type::Basic::Float ? 128 : 2;
                loadVar(ls, src, ir->src, 5);
                ls.push_back(new Arm::Store(src, mem));
            } else if (auto ir = irBase->As<IR::Call>()) {
                int argc = (int)ir->args.size();
                for (int i = 0; i < std::min(argc, 4); ++i) {
                    loadVar(ls, i, ir->args[i], 5);
                }
                for (int i = argc - 1; i >= 4; --i) {
                    loadVar(ls, 4, ir->args[i], 5);
                    ls.push_back(new Arm::Push(4));
                }
                ls.push_back(new Arm::Call(ir->func->name));
                addVar(ls, Register::sp, Register::sp,
                       4 * std::max(0, argc - 3), 1);
                if (ir->dest) {
                    storeVar(ls, 0, ir->dest, 5);
                }
            }

            // Terminator
            else if (auto ir = irBase->As<IR::Return>()) {
                loadVar(ls, 0, ir->ret, 5);
                blk->Connect(armFunc->exit);
            } else if (auto ir = irBase->As<IR::Jump>()) {
                blk->Connect(blocks[ir->nxt]);
            } else if (auto ir = irBase->As<IR::Branch>()) {
                if (ir->lhs->type.isPtr ||
                    ir->lhs->type.base == Type::Basic::Int) {
                    loadVar(ls, 0, ir->lhs, 5);
                    loadVar(ls, 1, ir->rhs, 5);
                    ls.push_back(new Arm::Cmp(0, 1));
                } else {
                    loadVar(ls, 128, ir->lhs, 5);
                    loadVar(ls, 129, ir->rhs, 5);
                    ls.push_back(new Arm::Cmp(128, 129));
                }
                blk->Connect(ir->op, blocks[ir->t], blocks[ir->f]);
            }
        }
    };
    for (auto [ir, arm] : blocks) {
        genBlock(arm, arm->instList, ir);
    }

    armFunc->entry->AddInst(new Arm::Push(
        std::vector<Register>{Register::lr, Register::fp, 3, 2, 1, 0}));
    armFunc->entry->AddInst(new Arm::Mov(Register::fp, Register::sp));
    addVar(armFunc->entry->instList, Register::sp, Register::sp, -stackSize, 0);

    addVar(armFunc->exit->instList, Register::sp, Register::sp, stackSize + 16,
           0);
    armFunc->exit->AddInst(
        new Arm::Pop(std::vector<Register>{Register::pc, Register::fp}));
    armFunc->exit->AddInst(new Arm::Ret());

    armFunc->entry->Connect(blocks[func->entry]);
    blocks[func->exit]->Connect(armFunc->exit);
    return armFunc;
}

void* ArmGenerator::Run(void* input) {
    auto module = (IRModule*)input;
    auto program = new ArmProgram();
    for (auto [var, sym] : module->globalVars) {
        auto armvar = Arm::GlobalVar();
        armvar.name = var->name;
        armvar.size = sym->allocMem;
        auto init = new std::unordered_map<int, int>();
        if (auto info = sym->info->As<Semantics::LValue>()) {
            (*init)[0] = info->init.v.i;
        } else if (auto info = sym->info->As<Semantics::Pointer>()) {
            if (info->init != nullptr) {
                for (auto [pos, val] : *info->init) {
                    (*init)[pos] = val.v.i;
                }
            }
        }
        armvar.init = init;
        program->globalVars.push_back(armvar);
    }
    for (auto func : module->globalFuncs) {
        auto armfunc = Build(func);
        program->funcs.push_back(armfunc);
    }
    return program;
}
